package com.yuyu.tt.config;

import com.yuyu.tt.context.TenantContextHolder;
import org.springframework.core.task.TaskDecorator;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;

public class RequestHeritableContextDecorator implements TaskDecorator {
    @Override
    public Runnable decorate(Runnable r) {
        // 获取ProjectCode 显式的传递给子线程
        RequestAttributes ra = RequestContextHolder.currentRequestAttributes();
        ServletRequestAttributes sra = (ServletRequestAttributes) ra;
        HttpServletRequest request = sra.getRequest();
        String projectCode = request.getHeader("projectCode");

        return () -> {
            try {
                TenantContextHolder.setProjectCode(projectCode);
                r.run();
            } finally {
                TenantContextHolder.removeProjectCode();
            }
        };
    }
}
